{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "461d8d77",
   "metadata": {},
   "source": [
    "# Chapter 3 - Optimal Flows (Python Code)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85fdcd6",
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [],
   "source": [
    "! pip install --upgrade quantecon_book_networks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3d5fcb1",
   "metadata": {},
   "source": [
    "We begin with some imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a00ff7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import quantecon as qe\n",
    "import quantecon_book_networks\n",
    "import quantecon_book_networks.input_output as qbn_io\n",
    "import quantecon_book_networks.plotting as qbn_plt\n",
    "import quantecon_book_networks.data as qbn_data\n",
    "ch3_data = qbn_data.optimal_flows()\n",
    "export_figures = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ed711ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.optimize import linprog\n",
    "import networkx as nx\n",
    "import ot\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import cm\n",
    "from matplotlib.patches import Polygon\n",
    "from matplotlib.artist import Artist  \n",
    "quantecon_book_networks.config(\"matplotlib\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a5e9cd1",
   "metadata": {},
   "source": [
    "## Linear Programming and Duality\n",
    "\n",
    "### Betweenness centrality (by color and node size) for the Florentine families\n",
    "\n",
    "We load the Florentine Families data from the NetworkX package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e14e416",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = nx.florentine_families_graph()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5199cc62",
   "metadata": {},
   "source": [
    "Next we calculate betweenness centrality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dab1b1d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "bc_dict = nx.betweenness_centrality(G)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "700349de",
   "metadata": {},
   "source": [
    "And we produce the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33557ff6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(9.4, 9.4))\n",
    "\n",
    "plt.axis(\"off\")\n",
    "nx.draw_networkx(\n",
    "    G, \n",
    "    ax=ax,\n",
    "    pos=nx.spring_layout(G, seed=1234), \n",
    "    with_labels=True,\n",
    "    alpha=.8,\n",
    "    arrowsize=15,\n",
    "    connectionstyle=\"arc3,rad=0.1\",\n",
    "    node_size=[10_000*(size+0.1) for size in bc_dict.values()], \n",
    "    node_color=[cm.plasma(bc+0.4) for bc in bc_dict.values()],\n",
    ")\n",
    "if export_figures:\n",
    "    plt.savefig(\"figures/betweenness_centrality_1.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fa794a1",
   "metadata": {},
   "source": [
    "### Revenue maximizing quantities and a Python implementation of linear programming\n",
    "\n",
    "First we specify our linear program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59dd01db",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = ((2, 5),\n",
    "     (4, 2))\n",
    "b = (30, 20)\n",
    "c = (-3, -4) # minus in order to minimize"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f64d9118",
   "metadata": {},
   "source": [
    "And now we use SciPy's linear programing module to solve our linear program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0bf9ba9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.optimize import linprog\n",
    "result = linprog(c, A_ub=A, b_ub=b)\n",
    "print(result.x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c7bfbb",
   "metadata": {},
   "source": [
    "Here we produce a visualization of what is being done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41077bbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4.5))\n",
    "plt.rcParams['font.size'] = '14'\n",
    "\n",
    "# Draw constraint lines\n",
    "\n",
    "ax.plot(np.linspace(-1, 17.5, 100), 6-0.4*np.linspace(-1, 17.5, 100))\n",
    "ax.plot(np.linspace(-1, 5.5, 100), 10-2*np.linspace(-1, 5.5, 100))\n",
    "ax.text(10, 2.5, \"$2q_1 + 5q_2 \\leq 30$\")\n",
    "ax.text(1.5, 8, \"$4q_1 + 2q_2 \\leq 20$\")\n",
    "ax.text(-2, 2, \"$q_2 \\geq 0$\")\n",
    "ax.text(2.5, -0.7, \"$q_1 \\geq 0$\")\n",
    "\n",
    "# Draw the feasible region\n",
    "feasible_set = Polygon(np.array([[0, 0], \n",
    "                                 [0, 6], \n",
    "                                 [2.5, 5], \n",
    "                                 [5, 0]]))\n",
    "ax.add_artist(feasible_set)\n",
    "Artist.set_alpha(feasible_set, 0.2) \n",
    "\n",
    "# Draw the objective function\n",
    "ax.plot(np.linspace(-1, 5.5, 100), 3.875-0.75*np.linspace(-1, 5.5, 100), 'g-')\n",
    "ax.plot(np.linspace(-1, 5.5, 100), 5.375-0.75*np.linspace(-1, 5.5, 100), 'g-')\n",
    "ax.plot(np.linspace(-1, 5.5, 100), 6.875-0.75*np.linspace(-1, 5.5, 100), 'g-')\n",
    "ax.text(5.8, 1, \"revenue $ = 3q_1 + 4q_2$\")\n",
    "\n",
    "# Draw the optimal solution\n",
    "ax.plot(2.5, 5, \"o\", color=\"black\")\n",
    "ax.text(2.7, 5.2, \"optimal solution\")\n",
    "\n",
    "for spine in ['right', 'top']:\n",
    "    ax.spines[spine].set_color('none')\n",
    "    \n",
    "ax.set_xticks(())\n",
    "ax.set_yticks(())\n",
    "\n",
    "for spine in ['left', 'bottom']:\n",
    "    ax.spines[spine].set_position('zero')\n",
    "    \n",
    "ax.set_ylim(-1, 8)\n",
    "\n",
    "if export_figures:\n",
    "    plt.savefig(\"figures/linear_programming_1.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1c4f25b",
   "metadata": {},
   "source": [
    "## Optimal Transport\n",
    " \n",
    "\n",
    "### Transforming one distribution into another\n",
    "\n",
    "Below we provide code to produce a visualization of transforming one\n",
    "distribution into another in one dimension."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92b7150b",
   "metadata": {},
   "outputs": [],
   "source": [
    "σ = 0.1\n",
    "\n",
    "def ϕ(z):\n",
    "    return (1 / np.sqrt(2 * σ**2 * np.pi)) * np.exp(-z**2 / (2 * σ**2))\n",
    "\n",
    "def v(x, a=0.4, b=0.6, s=1.0, t=1.4):\n",
    "    return a * ϕ(x - s) + b * ϕ(x - t)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 4))\n",
    "\n",
    "x = np.linspace(0.2, 4, 1000)\n",
    "ax.plot(x, v(x), label=\"$\\\\phi$\")\n",
    "ax.plot(x, v(x, s=3.0, t=3.3, a=0.6), label=\"$\\\\psi$\")\n",
    "\n",
    "ax.legend(loc='upper left', fontsize=12, frameon=False)\n",
    "\n",
    "ax.arrow(1.8, 1.6, 0.8, 0.0, width=0.01, head_width=0.08)\n",
    "ax.annotate('transform', xy=(1.9, 1.9), fontsize=12)\n",
    "if export_figures:\n",
    "    plt.savefig(\"figures/ot_figs_1.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7464f721",
   "metadata": {},
   "source": [
    "### Function to solve a transport problem via linear programming\n",
    "\n",
    "Here we define a function to solve optimal transport problems using linear programming."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94e547a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def ot_solver(phi, psi, c, method='highs-ipm'):\n",
    "    \"\"\"\n",
    "    Solve the OT problem associated with distributions phi, psi\n",
    "    and cost matrix c.\n",
    "    Parameters\n",
    "    ----------\n",
    "    phi : 1-D array\n",
    "    Distribution over the source locations.\n",
    "    psi : 1-D array\n",
    "    Distribution over the target locations.\n",
    "    c : 2-D array\n",
    "    Cost matrix.\n",
    "    \"\"\"\n",
    "    n, m = len(phi), len(psi)\n",
    "\n",
    "    # vectorize c\n",
    "    c_vec = c.reshape((m * n, 1), order='F')\n",
    "\n",
    "    # Construct A and b\n",
    "    A1 = np.kron(np.ones((1, m)), np.identity(n))\n",
    "    A2 = np.kron(np.identity(m), np.ones((1, n)))\n",
    "    A = np.vstack((A1, A2))\n",
    "    b = np.hstack((phi, psi))\n",
    "\n",
    "    # Call sover\n",
    "    res = linprog(c_vec, A_eq=A, b_eq=b, method=method)\n",
    "\n",
    "    # Invert the vec operation to get the solution as a matrix\n",
    "    pi = res.x.reshape((n, m), order='F')\n",
    "    return pi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfdeb358",
   "metadata": {},
   "source": [
    "Now we can set up a simple optimal transport problem."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54e87d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi = np.array((0.5, 0.5))\n",
    "psi = np.array((1, 0))\n",
    "c = np.ones((2, 2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b81830ff",
   "metadata": {},
   "source": [
    "Next we solve using the above function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e044c8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "ot_solver(phi, psi, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44e6ab3",
   "metadata": {},
   "source": [
    "We see we get the same result as when using the Python optimal transport\n",
    "package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82d43df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "ot.emd(phi, psi, c) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "642f85ba",
   "metadata": {},
   "source": [
    "### An optimal transport problem solved by linear programming\n",
    "\n",
    "Here we demonstrate a more detailed optimal transport problem. We begin by defining a node class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9bcd7819",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Node:\n",
    "\n",
    "    def __init__(self, x, y, mass, group, name):\n",
    "\n",
    "        self.x, self.y = x, y\n",
    "        self.mass, self.group = mass, group\n",
    "        self.name = name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad91c7e4",
   "metadata": {},
   "source": [
    "Now we define a function for randomly generating nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d1bfd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import betabinom\n",
    "\n",
    "def build_nodes_of_one_type(group='phi', n=100, seed=123):\n",
    "\n",
    "    nodes = []\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    for i in range(n):\n",
    "        \n",
    "        if group == 'phi':\n",
    "            m = 1/n\n",
    "            x = np.random.uniform(-2, 2)\n",
    "            y = np.random.uniform(-2, 2)\n",
    "        else:\n",
    "            m = betabinom.pmf(i, n-1, 2, 2)\n",
    "            x = 0.6 * np.random.uniform(-1.5, 1.5)\n",
    "            y = 0.6 * np.random.uniform(-1.5, 1.5)\n",
    "            \n",
    "        name = group + str(i)\n",
    "        nodes.append(Node(x, y, m, group, name))\n",
    "\n",
    "    return nodes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45b701ab",
   "metadata": {},
   "source": [
    "We now generate our source and target nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efa68979",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_phi = 32\n",
    "n_psi = 32\n",
    "\n",
    "phi_list = build_nodes_of_one_type(group='phi', n=n_phi)\n",
    "psi_list = build_nodes_of_one_type(group='psi', n=n_psi)\n",
    "\n",
    "phi_probs = [phi.mass for phi in phi_list]\n",
    "psi_probs = [psi.mass for psi in psi_list]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61f09811",
   "metadata": {},
   "source": [
    "Now we define our transport costs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38af4ff7",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = np.empty((n_phi, n_psi))\n",
    "for i in range(n_phi):\n",
    "    for j in range(n_psi):\n",
    "        x0, y0 = phi_list[i].x, phi_list[i].y\n",
    "        x1, y1 = psi_list[j].x, psi_list[j].y\n",
    "        c[i, j] = np.sqrt((x0-x1)**2 + (y0-y1)**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cbe78a44",
   "metadata": {},
   "source": [
    "We solve our optimal transport problem using the Python optimal transport package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8cbef04",
   "metadata": {},
   "outputs": [],
   "source": [
    "pi = ot.emd(phi_probs, psi_probs, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fd92511",
   "metadata": {},
   "source": [
    "Finally we produce a graph of our sources, targets, and optimal transport plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa85e02e",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = nx.DiGraph()\n",
    "g.add_nodes_from([phi.name for phi in phi_list])\n",
    "g.add_nodes_from([psi.name for psi in psi_list])\n",
    "\n",
    "for i in range(n_phi):\n",
    "    for j in range(n_psi):\n",
    "        if pi[i, j] > 0:\n",
    "            g.add_edge(phi_list[i].name, psi_list[j].name, weight=pi[i, j])\n",
    "\n",
    "node_pos_dict={}\n",
    "for phi in phi_list:\n",
    "    node_pos_dict[phi.name] = (phi.x, phi.y)\n",
    "for psi in psi_list:\n",
    "    node_pos_dict[psi.name] = (psi.x, psi.y)\n",
    "\n",
    "node_color_list = []\n",
    "node_size_list = []\n",
    "scale = 8_000\n",
    "for phi in phi_list:\n",
    "    node_color_list.append('blue')\n",
    "    node_size_list.append(phi.mass * scale)\n",
    "for psi in psi_list:\n",
    "    node_color_list.append('red')\n",
    "    node_size_list.append(psi.mass * scale)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(7, 10))\n",
    "plt.axis('off')\n",
    "\n",
    "nx.draw_networkx_nodes(g, \n",
    "                       node_pos_dict, \n",
    "                       node_color=node_color_list,\n",
    "                       node_size=node_size_list,\n",
    "                       edgecolors='grey',\n",
    "                       linewidths=1,\n",
    "                       alpha=0.5,\n",
    "                       ax=ax)\n",
    "\n",
    "nx.draw_networkx_edges(g, \n",
    "                       node_pos_dict, \n",
    "                       arrows=True,\n",
    "                       connectionstyle='arc3,rad=0.1',\n",
    "                       alpha=0.6)\n",
    "if export_figures:\n",
    "    plt.savefig(\"figures/ot_large_scale_1.pdf\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bd33217",
   "metadata": {},
   "source": [
    "### Solving linear assignment as an optimal transport problem\n",
    "\n",
    "Here we set up a linear assignment problem (matching $n$ workers to $n$ jobs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f26dc4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 4\n",
    "phi = np.ones(n)\n",
    "psi = np.ones(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "756aa046",
   "metadata": {},
   "source": [
    "We generate our cost matrix (the cost of training the $i$th worker for the $j$th job)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "437274b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "c = np.random.uniform(size=(n, n))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6a66222",
   "metadata": {},
   "source": [
    "Finally, we solve our linear assignment problem as a special case of optimal transport."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69565581",
   "metadata": {},
   "outputs": [],
   "source": [
    "ot.emd(phi, psi, c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef9fc68a",
   "metadata": {},
   "source": [
    "### Python Spatial Analysis library\n",
    "\n",
    "Readers interested in computational optimal transport should also consider\n",
    "PySAL, the [Python Spatial Analysis library](https://pysal.org/). See, for\n",
    "example, https://pysal.org/spaghetti/notebooks/transportation-problem.html.\n",
    "\n",
    "### The General Flow Problem\n",
    "\n",
    "Here we solve a simple network flow problem as a linear program. We begin by\n",
    "defining the node-edge incidence matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e7fc46b",
   "metadata": {},
   "outputs": [],
   "source": [
    "A = (\n",
    "( 1,  1,  0,  0),\n",
    "(-1,  0,  1,  0),\n",
    "( 0,  0, -1,  1),\n",
    "( 0, -1,  0, -1)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23fac21d",
   "metadata": {},
   "source": [
    "Now we define exogenous supply and transport costs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8c5644d",
   "metadata": {},
   "outputs": [],
   "source": [
    "b = (10, 0, 0, -10)\n",
    "c = (1, 4, 1, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3fe351ee",
   "metadata": {},
   "source": [
    "Finally we solve as a linear program."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a257b5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = linprog(c, A_eq=A, b_eq=b, method='highs-ipm')\n",
    "print(result.x)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
